package com.liuqi.aiagent.service.impl;

import com.alibaba.fastjson.JSONObject;
import com.liuqi.aiagent.mapper.ModelMapper;
import com.liuqi.aiagent.mapper.WebhookToolMapper;
import com.liuqi.aiagent.service.AgentChatManager;
import com.liuqi.aiagent.service.AgentService;
import com.liuqi.aiagent.transfer.common.LLMModelParameter;
import com.liuqi.aiagent.transfer.entity.Model;
import com.liuqi.aiagent.transfer.entity.WebhookTool;
import com.liuqi.aiagent.transfer.vo.AgentInfoVo;
import com.liuqi.openai.chain.OpenAiChain;
import com.liuqi.openai.chain.impl.SimpleOpenAiChain;
import com.liuqi.openai.core.client.OpenAiClient;
import com.liuqi.openai.core.client.httpclient.OpenAiHttpClient;
import com.liuqi.openai.core.chat.Function;
import com.liuqi.openai.core.chat.Tool;
import com.liuqi.openai.data.message.SystemMessage;
import com.liuqi.openai.memory.chat.MessageWindowChatMemory;
import com.liuqi.openai.model.OpenAiChatModel;
import com.liuqi.openai.tool.webhook.WebHook;
import com.liuqi.openai.tool.webhook.WebHookTool;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import static com.liuqi.aiagent.util.StringUtil.notHasText;

/**
 * @author liuqi
 * @date 2025/7/18
 **/
@Service
@Slf4j
@AllArgsConstructor
public class AgentChatManagerImpl implements AgentChatManager {

    private final AgentService agentService;

    private final ModelMapper modelMapper;

    private final WebhookToolMapper webhookToolMapper;

    /**
     * 缓存key的拼接符号
     */
    private final String KEY_MONTAGE_SYMBOL = "@";

    /**
     * OpenAiChain 缓存器, 缓存数据结构
     * <p>
     * agent_name: {
     * username: { AgentCache }
     * }
     * </p>
     * <p>
     * 例如:
     * <p>
     * "Test-Agent": {
     * "admin": { AgentCache },
     * "user": { AgentCache }
     * },
     * "Ai-Agent": {
     * "admin": { AgentCache },
     * "user": { AgentCache }
     * }
     * </p>
     */
    private final ConcurrentHashMap<String, ConcurrentHashMap<String, AgentCache>> openAiChainCache = new ConcurrentHashMap<>();

    /**
     * OpenAiClient 缓存器
     * <p>
     * key = url@apikey
     */
    private final ConcurrentHashMap<String, OpenAiClient> openAiClientCache = new ConcurrentHashMap<>();

    @Override
    public OpenAiChain<String> findOpenAiChain(String agentName, String username) {
        ConcurrentHashMap<String, AgentCache> agentCaches = openAiChainCache.get(agentName);
        if (agentCaches == null) {
            synchronized (this) {
                if ((agentCaches = openAiChainCache.get(agentName)) == null) {
                    agentCaches = new ConcurrentHashMap<>();
                    openAiChainCache.put(agentName, agentCaches);
                }
            }
        }

        AgentCache agentCache = agentCaches.get(username);
        if (agentCache == null) {
            synchronized (this) {
                if ((agentCache = agentCaches.get(username)) == null) {
                    // 创建 OpenAiChain
                    agentCache = createOpenAiChain(agentName, username);
                    agentCaches.put(username, agentCache);
                }
            }
        }

        return agentCache.openAiChain;
    }

    @Override
    public void removeCache(String agentName) {
        openAiChainCache.remove(agentName);
        log.info("移除缓存Agent, agentName = " + agentName);
    }

    @Override
    public Integer findAgentId(String agentName, String username) {
        AgentCache agentCache = getAgentCache(agentName, username);
        if (agentCache != null) {
            return Integer.parseInt(agentCache.agentId);
        }
        return null;
    }

    @Override
    public Integer findModelId(String agentName, String username) {
        AgentCache agentCache = getAgentCache(agentName, username);
        if (agentCache != null) {
            return agentCache.modelId;
        }
        return null;
    }

    @Override
    public String findModelName(String agentName, String username) {
        AgentCache agentCache = getAgentCache(agentName, username);
        if (agentCache != null) {
            return agentCache.modelName;
        }
        return null;
    }

    private AgentCache getAgentCache(String agentName, String username) {
        ConcurrentHashMap<String, AgentCache> agentCaches = openAiChainCache.get(agentName);
        if (agentCaches == null) {
            return null;
        }
        return agentCaches.get(username);
    }

    private String chainKey(String agentName, String username) {
        return agentName + KEY_MONTAGE_SYMBOL + username;
    }

    /**
     * 创建 OpenAiChain
     *
     * @param agentName
     * @param username
     * @return
     */
    private AgentCache createOpenAiChain(String agentName, String username) {
        log.info("AgentChatManager: 创建[OpenAiChain], agentName[{}], username[{}]", agentName, username);
        try {
            // 查询 Agent
            AgentInfoVo agentInfo = checkAgentInfoVo(agentService.info(agentName));
            if (log.isDebugEnabled()) {
                log.debug("agentInfo: " + JSONObject.toJSONString(agentInfo));
            }
            // 查询模型
            Model model = checkModel(modelMapper.findById(Integer.parseInt(agentInfo.getModelId())));
            if (log.isDebugEnabled()) {
                log.debug("model: " + model);
            }

            // open ai 客户端
            OpenAiClient openAiClient = getAndCreateOpenAiClient(model);

            // 大语言模型入参
            LLMModelParameter llmModelParameter = agentInfo.getLlmModelParameter();

            // 创建 OpenAI 对话模型
            OpenAiChatModel chatModel = OpenAiChatModel.builder()
                    .openAiClient(openAiClient)
                    .modelName(model.getModelName())
                    .temperature(llmModelParameter.getTemperature())
                    .topP(llmModelParameter.getTopP())
                    .presencePenalty(llmModelParameter.getPresencePenalty())
                    .frequencyPenalty(llmModelParameter.getFrequencyPenalty())
                    .maxTokens(llmModelParameter.getMaxTokens())
                    .enableThinking(llmModelParameter.getEnableThinking())
                    .thinkingBudget(llmModelParameter.getThinkingBudget())
                    .build();

            // 构建外部工具
            List<WebHookTool> webHookTools = null;
            List<String> webhookIds = agentInfo.getWebhookIds();
            if (webhookIds != null && !webhookIds.isEmpty()) {
                List<WebhookTool> tools = webhookToolMapper.findAllByIds(webhookIds.stream().map(Integer::parseInt).collect(Collectors.toList()));
                webHookTools = toAiTool(tools);
            }

            // 创建 SimpleOpenAiChain
            SimpleOpenAiChain openAiChain = SimpleOpenAiChain.builder()
                    .id(chainKey(agentName, username))
                    .chatModel(chatModel)
                    .chatMemory(
                            MessageWindowChatMemory.withMaxMessages(10,
                                    /* 设置提示词 */
                                    SystemMessage.from(agentInfo.getPrompt()))
                    )
                    .customTools(webHookTools)
                    .build();

            log.info("AgentChatManager: 创建OpenAiChain成功. agentName[{}], username[{}]", agentName, username);

            return AgentCache.builder()
                    .username(username)
                    .agentId(agentInfo.getId())
                    .agentName(agentInfo.getName())
                    .modelId(model.getId())
                    .modelName(model.getModelName())
                    .openAiChain(openAiChain)
                    .build();
        } catch (Exception e) {
            log.error("AgentChatManager: 创建OpenAiChain失败", e);
            throw e;
        }
    }

    /**
     * 从缓存中获取 OpenAiClient
     *
     * @param model
     * @return
     */
    private OpenAiClient getAndCreateOpenAiClient(Model model) {
        String key = model.getUrl() + KEY_MONTAGE_SYMBOL + model.getApiKey();
        OpenAiClient openAiClient = openAiClientCache.get(key);
        if (openAiClient != null) {
            return openAiClient;
        } else {
            synchronized (this) {
                if ((openAiClient = openAiClientCache.get(key)) != null) {
                    return openAiClient;
                } else {
                    log.info("创建OpenAiHttpClient, url[{}], apikey[{}]", model.getUrl(), model.getApiKey());
                    openAiClient = OpenAiHttpClient.builder()
                            .baseUrl(model.getUrl())
                            .openAiApiKey(model.getApiKey())
                            .build();
                    openAiClientCache.put(key, openAiClient);
                    return openAiClient;
                }
            }
        }
    }

    private AgentInfoVo checkAgentInfoVo(AgentInfoVo agentInfo) {
        if (agentInfo == null) {
            throw new NullPointerException("agent 不存在");
        }
        if (agentInfo.getModelId() == null) {
            throw new NullPointerException("模型为空");
        }
        if (agentInfo.getPrompt() == null) {
            throw new NullPointerException("提示词为空");
        }
        return agentInfo;
    }

    private Model checkModel(Model model) {
        if (model == null) {
            throw new NullPointerException("查询模型不存在");
        }
        if (notHasText(model.getModelName())) {
            throw new NullPointerException("模型名称为空");
        }
        if (notHasText(model.getUrl())) {
            throw new NullPointerException("模型服务地址为空");
        }
        if (notHasText(model.getApiKey())) {
            throw new NullPointerException("模型ApiKey为空");
        }
        return model;
    }

    private List<WebHookTool> toAiTool(List<WebhookTool> tools) {
        if (tools.isEmpty()) {
            return null;
        }
        List<WebHookTool> aiTools = new ArrayList<>(tools.size());
        tools.forEach(tool -> {
            WebHook webHook = WebHook.builder()
                    .servers(JSONObject.parseArray(tool.getServers(), String.class))
                    .headers(JSONObject.parseObject(tool.getHeaders(), Map.class))
                    .contentType(tool.getContentType())
                    .path(tool.getPath())
                    .build();
            Function function = JSONObject.parseObject(tool.getFunction(), Function.class);
            WebHookTool webHookTool = WebHookTool.builder()
                    .webHook(webHook)
                    .toolName(function.getName())
                    .chatTool(new Tool("function", function))
                    .build();
            aiTools.add(webHookTool);
        });
        return aiTools;
    }

    /**
     * Agent 缓存消信息
     */
    @Builder
    private static class AgentCache {
        private final String username;
        private final String agentId;
        private final String agentName;
        private final Integer modelId;
        private final String modelName;
        private final OpenAiChain<String> openAiChain;
    }

}
